import json
import os
import shapely.geometry as geom
import numpy as np

obj_list = ["WHITE_LABEL"]

if __name__ == "__main__":
    obj_dir = r"C:\Users\Administrator\Desktop\yolo_4dfl\datasets\smart_label_data"
    name_list = os.listdir(obj_dir)
    for name1 in name_list:
        if name1.split(".")[1] == "json":
            img_path = os.path.join(obj_dir, name1.split(".")[0] + ".png")
            json_path = os.path.join(obj_dir, name1.split(".")[0] + ".json")
            txt_path = os.path.join(obj_dir, name1.split(".")[0] + ".txt")

            row_list = []
            with open(json_path, "r") as f:
                dict1 = json.load(f)
                imageWidth = dict1["imageWidth"]
                imageHeight = dict1["imageHeight"]
                img_name = dict1["imagePath"]
                shapes = dict1["shapes"]
                for shape in shapes:
                    if shape["label"] in obj_list:
                        x1 = shape["points"][0][0] / imageWidth
                        y1 = shape["points"][0][1] / imageHeight
                        x2 = shape["points"][1][0] / imageWidth
                        y2 = shape["points"][1][1] / imageHeight
                        x3 = shape["points"][2][0] / imageWidth
                        y3 = shape["points"][2][1] / imageHeight
                        x4 = shape["points"][3][0] / imageWidth
                        y4 = shape["points"][3][1] / imageHeight

                        polygon = geom.Polygon([(x1, y1), (x2, y2), (x3, y3), (x4, y4)])
                        rectangle = geom.Polygon([(0, 0), (1, 0), (1, 1), (0, 1)])
                        # 求交集
                        intersection = polygon.intersection(rectangle)
                        # 计算交集面积占总面积的比例
                        intersection_ratio = intersection.area / polygon.area
                        if intersection_ratio < 0.4:
                            continue

                        # 获取交集的 bounding box
                        minx, miny, maxx, maxy = intersection.bounds
                        bounding_box = [(minx, miny), (maxx, maxy)]
                        print("Bounding Box:", bounding_box)

                        if not ("nan" in str(bounding_box)):
                            X, Y = intersection.exterior.coords.xy
                            if len(X) != 5:
                                continue
                            coords = np.array(
                                [
                                    [X[0], Y[0]],
                                    [X[1], Y[1]],
                                    [X[2], Y[2]],
                                    [X[3], Y[3]],
                                ]
                            )
                            index1 = np.argmin(np.sum(coords, axis=1))
                            index3 = np.argmax(np.sum(coords, axis=1))
                            x1y1 = coords[index1]
                            x3y3 = coords[index3]

                            coords = np.delete(coords, [index1, index3], axis=0)
                            index2 = np.argmin(np.sum(coords * np.array([-1, 1]), axis=1))
                            index4 = np.argmax(np.sum(coords * np.array([-1, 1]), axis=1))
                            x2y2 = coords[index2]
                            x4y4 = coords[index4]
                            row1 = [
                                0,
                                x1y1[0],
                                x1y1[1],
                                x2y2[0],
                                x2y2[1],
                                x3y3[0],
                                x3y3[1],
                                x4y4[0],
                                x4y4[1],
                            ]
                            if ((maxy - miny) * (maxx - minx)) < 0.003:
                                continue
                            row1 = list(map("{:.6f}".format, row1))
                            row1[0] = 0
                            row_list.append(row1)

            with open(txt_path, "w", newline="") as f:
                for row in row_list:
                    for data in row:
                        f.write(str(data))
                        f.write(" ")
                    f.write("\r\n")
